import pandas as pd
import numpy as np
from sklearn.ensemble import RandomForestRegressor
from sklearn.metrics import r2_score, mean_absolute_error, mean_squared_error
import joblib

from data_graph import train_data_keys

# ========= 特征列 & 标签列 =========
feature_cols = [
    "driver_output_cap", "total_sink_cap",
    "driver_to_sink_dx", "driver_to_sink_dy",
    "sink_x", "sink_y", "max_driver_input_slew"
]
target_col = "target_delay"

# ========= 加载数据集 =========
df = pd.read_csv("rf_feature.csv")

# ========= 训练/测试划分 =========
train_df = df[df["design_name"].isin(train_data_keys)].reset_index(drop=True)
test_df = df[~df["design_name"].isin(train_data_keys)].reset_index(drop=True)

X_train = train_df[feature_cols]
y_train = train_df[target_col]

X_test = test_df[feature_cols]
y_test = test_df[target_col]

print(f" Train samples: {len(train_df)}")
print(f" Test samples : {len(test_df)}")

# ========= 模型训练 =========
model = RandomForestRegressor(
    n_estimators=500,
    max_depth=25,
    # min_samples_leaf=10,
    # min_samples_split=15,
    #max_features=3,
    #max_samples=0.8,
    #bootstrap=True,
    # random_state=42,
    n_jobs=-1,
)


model.fit(X_train, y_train)

# ========= 测试预测 =========
y_pred = model.predict(X_test)

# ========= 总体评估 =========
r2 = r2_score(y_test, y_pred)
mae = mean_absolute_error(y_test, y_pred)
rmse = np.sqrt(mean_squared_error(y_test, y_pred))

print("\n Overall Evaluation on Test Set:")
print(f"R² Score       : {r2:.4f}")
print(f"MAE (ns)       : {mae:.4f}")
print(f"RMSE (ns)      : {rmse:.4f}")

# ========= 每个 design 的 R² =========
def print_r2_by_design(model, df, feature_cols, target_col, split_name):
    design_groups = df.groupby("design_name")
    design_r2 = {}

    for design, group in design_groups:
        X = group[feature_cols]
        y_true = group[target_col]
        y_pred = model.predict(X)
        r2 = r2_score(y_true, y_pred)
        design_r2[design] = r2

    print(f"\n R² per design in {split_name}:")
    for design, score in sorted(design_r2.items(), key=lambda x: x[0]):
        print(f"  {design:<20}  {score:.4f}")

# 输出每个 design 的 R²
print_r2_by_design(model, train_df, feature_cols, target_col, split_name="Train Set")
print_r2_by_design(model, test_df, feature_cols, target_col, split_name="Test Set")

# === 保存模型（可选）===

# joblib.dump(model, "rf_net_delay_model.pkl")
# print("模型已保存为 rf_net_delay_model.pkl")
